import os
import mplfinance as mpf
import pandas as pd
import numpy as np
from sqlalchemy import create_engine
from matplotlib.figure import Figure
from matplotlib import gridspec
from pred import auto_pred


def get_k_line(d: dict):
    db_future_main = 'mysql+pymysql://hf_user:290202@192.168.3.179:3306/future?charset=utf8&use_unicode=1'
    engine_future_main = create_engine(db_future_main)
    choose = pd.read_sql(f"select * from {d['dbname']}", engine_future_main)  # 从future数据库读取主力合约
    return choose


def plot_kline(d: dict, fig: Figure):
    dir_ = f"data/{d['name']}"
    if not os.path.exists(dir_):
        os.makedirs(dir_)

    code_tick = d['dbname'].replace('_', '.')

    path = f"{dir_}/{code_tick}_choose_df.xlsx"
    bool_data_path = f"cache/{code_tick}_boll_df.xlsx"
    if os.path.exists(bool_data_path):
        bool_data = pd.read_excel(bool_data_path, index_col=0)
    else:
        auto_pred(code_tick)  # 调用自动预测函数，自动生成未来的预测并保存本地
        bool_data = pd.read_excel(bool_data_path, index_col=0)

    if not os.path.exists(path):
        db_future_main = 'mysql+pymysql://hf_user:290202@192.168.3.179:3306/future?charset=utf8&use_unicode=1'
        engine_future_main = create_engine(db_future_main)

        choose = pd.read_sql(f"select * from {d['dbname']}", engine_future_main)  # 从future数据库读取主力合约
        choose.index = choose['trade_date']
        choose.index = pd.to_datetime(choose.index)
        choose.index.name = 'Date'
        choose = choose.sort_index(ascending=True)
        choose = choose.iloc[-90:, :]
        choose["Open"] = choose["open"]
        choose["High"] = choose["high"]
        choose["Low"] = choose["low"]
        choose["Close"] = choose["close"]
        choose["Volume"] = choose["vol"]
        choose = choose.loc[:, ["Open", "High", "Low", "Close", "Volume"]]
        choose.to_excel(f"{dir_}/{code_tick}_choose_df.xlsx")
    else:
        choose = pd.read_excel(f"{dir_}/{code_tick}_choose_df.xlsx", index_col=0)

    y_history_max = np.max(choose.loc[:, ["Open", "High", "Low", "Close"]].values)
    y_history_min = np.min(choose.loc[:, ["Open", "High", "Low", "Close"]].values)

    spec = gridspec.GridSpec(ncols=6, nrows=4, figure=fig)
    ax1 = fig.add_subplot(spec[:3, :5])
    ax2 = fig.add_subplot(spec[-1:, :5])

    mpf.plot(choose, ax=ax1, volume=ax2, style='yahoo', type='candle')  # 调用mpf绘图方法，自己写太麻烦
    ax1.yaxis.set_ticks_position('left')  # 原始图的标签和刻度在右边，因为右边还有预测，所以放在左边

    ax3 = fig.add_subplot(spec[:3, 5:])

    bool_data.index = list(range(len(bool_data)))
    y_pred_max = np.max(bool_data.values)
    y_pred_min = np.min(bool_data.values)

    # 计算y_max和y_min是为了两个图的上下幅度一致
    y_max = max(y_history_max, y_pred_max)
    y_min = min(y_history_min, y_pred_min)

    ax3.plot(bool_data['v_50'], color='black', ls='--')
    v_25 = bool_data['v_25'].to_list()  # 百分之25分位数的预测
    v_75 = bool_data['v_75'].to_list()  # 百分之75分位数的预测

    ax1.set_ylim(y_min, y_max)
    ax3.set_ylim(y_min, y_max)

    ax3.fill_between(bool_data['v_50'].index.to_list(), y1=v_25, y2=v_75, color='gray')  # 绘制阴影

    for ax in [ax1, ax2]:  # 图例不旋转，mpf绘制的会自动旋转，不好看
        for tick in ax.get_xticklabels():
            tick.set_rotation(0)

    for ax in [ax1, ax2, ax3]:  # 右边和上边不显示，（显示效果不好看）
        ax.spines['top'].set_visible(False)
        ax.spines['right'].set_visible(False)

    # fig.savefig(path, bbox_inches="tight")

    return fig


if __name__ == "__main__":
    data = pd.read_excel(f"data/摘要.xlsx", index_col=0, sheet_name='焦煤')
    data = data[data.page == 1]
    data = data[data.save_name == "K线图"]
    d_ = data.iloc[0, :].to_dict()
    p = plot_kline(d_)
